import equinox as eqx
from jaxtyping import Array, Float, Int, Bool
import jax
import jax.numpy as jnp
import math
import numpy as np

class TemporalEnc(eqx.Module):
    pos_vec: Array
    dim: int
    mask: Array

    def __init__(self, dim: int):
        self.pos_vec = jnp.asarray([math.pow(10000.0, 2.0 * (i // 2) / dim) for i in range(dim)]) # (dim, )
        self.dim = dim

        mask = np.zeros(dim, dtype=bool)
        mask[0::2] = True
        self.mask = jnp.asarray(mask)

    @eqx.filter_jit
    def __call__(self, ts: Float[Array, "T"]):
        result = ts[:, None] / jax.lax.stop_gradient(self.pos_vec)[None, :] # (T, dim)
        result = jnp.where(self.mask[None, :], jnp.sin(result), jnp.cos(result))
        return result # (T, dim)

class ShiftedTemporalEnc(eqx.Module):
    position: Array
    div_term: Array
    layer_time_delta: eqx.Module

    def __init__(self, dim: int, max_len: int, key: Array):
        self.position = jnp.arange(0, max_len, dtype=float) # (max_len, )
        self.div_term = jnp.exp(jnp.arange(0, dim, 2, dtype=float)* -(math.log(10000.) / dim)) # (dim//2, )
        self.layer_time_delta = eqx.nn.Linear(1, dim // 2, False, key=key)

    @eqx.filter_jit
    def __call__(self, ts: Float[Array, "T"]):
        temp = jnp.concatenate((jnp.zeros(1), ts), 0)
        dts = temp[1:] - temp[:-1]
        dts = jnp.where(dts<0, 0., dts)
        phi = jax.vmap(jax.lax.stop_gradient(self.layer_time_delta))(dts[:, None]) # (T, dim//2)
        length = dts.shape[0]
        arc = jax.lax.stop_gradient(self.position)[:length, None] * jax.lax.stop_gradient(self.div_term) # (T, dim//2)
        temp = arc + phi
        pe_sin = jnp.sin(temp)
        pe_cos = jnp.cos(temp)
        pe = jnp.concatenate((pe_sin, pe_cos), 1)
        return pe # (T, dim)
